#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
FedCP 入口脚本（重写版）
支持：
    • 经典数据集：mnist / cifar …（保持向后兼容）
    • MIMIC-IHM & MIMIC-Phenotyping
    • 新模型：mult  text  tsmixed
"""
import os, sys, time, argparse, warnings, random, copy
import torch
import numpy as np

from flcore.servers.servercp import FedCP          # 服务器
from flcore.clients.fedcp_wrapper import wrap   # 模型包装器
from flcore.trainmodel.model import (
    MULTCrossModel, TextModel, TSMixed                    # 你上传的 3 个模型
)

warnings.filterwarnings("ignore")
torch.manual_seed(0)
np.random.seed(0)
random.seed(0)


# ----------  util ----------
def set_device(gpu_id: str):
    if torch.cuda.is_available() and gpu_id != "cpu":
        torch.cuda.set_device(int(gpu_id))
        return torch.device("cuda")
    return torch.device("cpu")


# ----------  main ----------
def main():
    parser = argparse.ArgumentParser()

    # ─── 数据集 ───
    parser.add_argument("-data", "--dataset", choices=["mnist", "cifar10", "ihm", "pheno"],
                        default="ihm")
    parser.add_argument("-droot", "--data_root",
                        default="./MIMICdataprocess/MultimodalMIMIC/Data/ihm")
    parser.add_argument("-nb", "--num_labels", type=int, default=2)
    parser.add_argument("--file_path", default="./MultimodalMIMIC/Data/ihm",
                        help="目录下要有 trainp2x_data.pkl 等文件")
    parser.add_argument("--train_batch_size", type=int, default=32)
    parser.add_argument("--eval_batch_size", type=int, default=32)
    parser.add_argument("--chunk", action="store_true")  # 是否加载 token_attn.pkl
    parser.add_argument("--debug", action="store_true")  # 仅取 100 条样本
    # ─── 客户端本地训练参数 (NEW) ───
    parser.add_argument("--local_learning_rate", "-lr", type=float, default=1e-3,
                        help="每个客户端本地优化器的学习率")
    parser.add_argument("--local_epochs", "-le", type=int, default=5,
                        help="每轮联邦迭代中，客户端本地训练的 epoch 数")
    parser.add_argument("-bs", "--batch_size", type=int, default=32,
                        help="本地 DataLoader 的 batch size")
    parser.add_argument("--local_steps", "-ls", type=int, default=100,
                        help="每轮联邦中，客户端的本地梯度更新步数")

    # ─── 模型与任务 ───
    parser.add_argument("-m", "--model",
                        choices=["cnn", "resnet18", "fastText", "mult", "text", "tsmixed"],
                        default="mult")
    parser.add_argument("--task", choices=["ihm", "pheno"], default="ihm")
    parser.add_argument("--modeltype", default="TS_Text",  # ← MULTCrossModel 要用
                        choices=["TS_Text", "TS", "Text"])
    parser.add_argument("--cross_method", default="MulT",
                        choices=["MulT", "self_cross", "MAGGate", "Outer"])
    parser.add_argument("--cross_layers", type=int, default=1)
    parser.add_argument("--bert_name",
                        default="/home/liuning/miaohr/bioformer-8L",
                        help="本地目录或 HuggingFace 名称")

    # 通用超参
    parser.add_argument("--embed_dim", type=int, default=128)
    parser.add_argument("--num_heads", type=int, default=4)
    parser.add_argument("--layers", type=int, default=2)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--kernel_size", type=int, default=3)
    parser.add_argument("--tt_max", type=int, default=48)

    # 额外旗标（MULT / TSMixed 会用到）
    parser.add_argument("--irregular_learn_emb_ts", action="store_true")
    parser.add_argument("--irregular_learn_emb_text", action="store_true")
    parser.add_argument("--reg_ts", action="store_true")  # 用规则 TS 通道
    parser.add_argument("--TS_mixup", action="store_true")
    parser.add_argument("--mixup_level", default="batch",
                        choices=["batch", "batch_seq", "batch_seq_feature"])
    parser.add_argument("--Interp", action="store_true")
    parser.add_argument("--TS_model", default="Atten",
                        choices=["Atten", "LSTM", "CNN"])
    parser.add_argument("--embed_time", type=int, default=16)
    parser.add_argument("--embed_dropout", type=float, default=0.1)
    parser.add_argument("--max_length", type=int, default=128,
                        help="TSNote_Irg 内部截断的最大长度")
    parser.add_argument("--pad_to_max_length", action="store_true",
                        help="若给定，则 TSNote_Irg 会把文本 pad 到 --max_length")
    parser.add_argument("--ratio_notes_order", type=float, default=None,
                        help="控制 note 顺序抽样的比例；留空(None)则按完整顺序读取")
    parser.add_argument("--num_of_notes", type=int, default=5,
                        help="每条 stay 最多取多少条临床笔记；≥1")
    parser.add_argument("--ts_dim", type=int, default=34,
                        help="单条时间序列的通道数 × 2 (含 mask)")

    # ─── 联邦学习 ───
    parser.add_argument("-gr", "--global_rounds", type=int, default=200)
    parser.add_argument("-nc", "--num_clients", type=int, default=3)
    parser.add_argument("-jr", "--join_ratio", type=float, default=0.5)
    parser.add_argument("--random_join_ratio", action="store_true")
    parser.add_argument("--eval_gap", type=int, default=1)
    # --- 数据划分 ---
    parser.add_argument("--seed", type=int, default=42,  # 随机种子
                        help="数据划分用的 RNG 种子")
    parser.add_argument("--split_strategy", default="iid",
                        choices=["iid", "dirichlet"],
                        help="客户端数据划分方式 (暂仅实现 iid)")
    parser.add_argument("--dirichlet_alpha", type=float, default=0.5,
                        help="Dirichlet 非 IID 的 α 值")

    # 运行次数 / 正则等旧旗标
    parser.add_argument("-t", "--times", type=int, default=1)
    parser.add_argument("-lam", "--lamda", type=float, default=0.0)

    # 设备
    parser.add_argument("-dev", "--device", choices=["cpu", "cuda"], default="cuda")
    parser.add_argument("-did", "--device_id", default="0")

    # 算法（FedCP、FedAvg…）
    parser.add_argument("-algo", "--algorithm", default="FedCP")

    args = parser.parse_args()

    # ───── 设备 ─────
    args.device = set_device(args.device_id if args.device=="cuda" else "cpu")
    print(f"▶ device = {args.device}")

    # ───── 模型构建 ─────
    if args.model in ("mult", "text", "tsmixed"):
        # 数据维度（可按自己数据修改）
        orig_d_ts, orig_reg_d_ts, ts_seq_num  = args.ts_dim, args.ts_dim, args.tt_max
        orig_d_txt, text_seq_num              = 768, args.tt_max

        if args.model == "mult":
            # ① 先加载你本地的 BioFormer-8L
            from transformers import AutoModel  # 顶部已 import 过可删
            print(f"▶ Loading encoder from {args.bert_name}")
            biobert = AutoModel.from_pretrained(args.bert_name).to(args.device)

            # ② 取到隐藏维度（BioFormer-8L 是 256）
            orig_d_txt = biobert.config.hidden_size

            # ③ 创建 MULTCrossModel，并把 biobert 传进去
            core = MULTCrossModel(
                args,
                device=args.device,
                orig_d_ts=orig_d_ts, orig_reg_d_ts=orig_reg_d_ts,
                orig_d_txt=orig_d_txt,
                ts_seq_num=ts_seq_num, text_seq_num=ts_seq_num,
                Biobert=biobert  # ★ 关键：注入已加载的编码器
            )
        elif args.model == "text":
            core = TextModel(args, device=args.device, orig_d_txt=orig_d_txt)
        else:  # tsmixed
            core = TSMixed(args, device=args.device,
                           orig_d_ts=orig_d_ts, orig_reg_d_ts=orig_reg_d_ts,
                           ts_seq_num=ts_seq_num)

        args.model = wrap(core)           # ✦✦ 关键：包装成 FedCP 兼容 ✦✦
        args.nb = args.num_labels         # servercp 里要用
    else:
        # 原有 cnn / resnet / fastText 逻辑保持不变
        from flcore.trainmodel.models import create_model
        args.model = wrap(create_model(args.model, args.num_labels).to(args.device))

    # ───── 其他路径信息（供 dataset 读取） ─────
    args.data_path = args.data_root
    args.batch_size = args.train_batch_size
    os.environ["FEDCP_DATA_ROOT"] = args.data_root
    # ───── 启动服务器 / 训练 ─────
    server = FedCP(args, times=int(time.time()))  # 用当前时间戳记一次实验
    server.train()


if __name__ == "__main__":
    t0 = time.time()
    main()
    print(f"✓ finished  |  total time = {time.time()-t0:.1f}s")
